from __future__ import absolute_import
import functools
import torch
from torch import nn
from torch.nn import functional as F
import torchvision
import random
import numpy as np

from .res_tgcn import res_tGCN
from .res_sgcn import res_sGCN
from .batch_res_sgcn import batch_res_sGCN
from .tgcn import tGCN
from .batch_tgcn import batch_tGCN
from .sgcn import sGCN
from .batch_sgcn import batch_sGCN



class gcn_model(nn.Module):
    #def __init__(self,seq_len=8,spatch_num=4,tpatch_num=4,graph_arch='gcn',gcn_num=4,g_feat=512,layer_num=4,g_pool='max'):
    def __init__(self,args):
        super(gcn_model, self).__init__()
        self.spatch_num=args.spatch_num
        self.tpatch_num=args.tpatch_num

        self.graph_arch=args.graph_arch
        self.feat_dim = 2048
        self.layer_num=args.layer_num
        self.tlayer_num=args.tlayer_num
        self.gcn_num=args.gcn_num
        self.g_feat_dim=args.g_feat
        self.seq_len=args.S
        self.g_pool=args.g_pool
        self.use_batch=args.use_batch
        self.args=args

        self.res_sgcn_net=batch_res_sGCN(in_feat_dim=self.feat_dim,out_feat_dim=self.feat_dim,layer_num=self.layer_num,g_pool=self.g_pool,args=args)
        self.tgcn_net=nn.ModuleList([batch_tGCN(in_feat_dim=self.feat_dim,out_feat_dim=self.g_feat_dim,layer_num=self.tlayer_num,g_pool=self.g_pool,args=args) for _ in range(self.gcn_num)])
    
        #初始化参数
        self._init_params()

        #初始化参数
    def _init_params(self):
        pass
         

    def forward(self, x):
       
        seq_len = self.seq_len #seq_len
        # 384,128对应 [batch_size*seq_len, self.feat_dim, 24, 8]
     
        s_graph_feature=self.res_sgcn_model(x,seq_len)
        t_graph_feature=self.tgcn_model(x,seq_len)
        graph_feature=[s_graph_feature,t_graph_feature]
    
        return  graph_feature
           
           
    def res_sgcn_model(self,x,seq_len):
            #GCN
         
            gx=x
            g_batch_size=gx.size(0) # g_batch_size=batch_size*seq_len
            batch=g_batch_size//seq_len # original batch size
            gh=gx.size(2)
            gw=gx.size(3)
            gx=gx.view(g_batch_size,self.feat_dim,self.spatch_num,int(gh/self.spatch_num),gw)
            #转换维度
            gx=gx.permute(0,2,1,3,4) # 互换 self.feat_dim和patch_num所在的维度，此时gx shape:[g_batch_size,patch_num,self.feat_dim,int(gh/patch_num),gw]
            gx=gx.contiguous() 
            gx=gx.view(g_batch_size*self.spatch_num,self.feat_dim,int(gh/self.spatch_num),gw) #shape:[batch_size*seq_len*patch_size,self.feat_dim,int(gh/patch_size),gw]
            #pooling
            gx = F.avg_pool2d(gx, gx.size()[2:]) # shape[batch_size*seq_len*patch_size,self.feat_dim,1,1]
            gx=gx.view(batch,seq_len*self.spatch_num,self.feat_dim)

            feature=self.res_sgcn_net(gx)
            return feature
          

    def tgcn_model(self,x,seq_len):
            #GCN
           
            gx=x
            g_batch_size=gx.size(0) # g_batch_size=batch_size*seq_len
            batch=g_batch_size//seq_len # original batch size
            gh=gx.size(2)
            gw=gx.size(3)
            gx=gx.view(g_batch_size,self.feat_dim,self.tpatch_num,int(gh/self.tpatch_num),gw)
            #转换维度
            gx=gx.permute(0,2,1,3,4) # 互换 self.feat_dim和patch_num所在的维度，此时gx shape:[g_batch_size,patch_num,self.feat_dim,int(gh/patch_num),gw]
            gx=gx.contiguous() 
            gx=gx.view(g_batch_size*self.tpatch_num,self.feat_dim,int(gh/self.tpatch_num),gw) #shape:[batch_size*seq_len*patch_size,self.feat_dim,int(gh/patch_size),gw]
            #pooling
            gx = F.avg_pool2d(gx, gx.size()[2:]) # shape[batch_size*seq_len*patch_size,self.feat_dim,1,1]
            gx=gx.view(batch,seq_len,self.tpatch_num,self.feat_dim)


            if(self.use_batch):
                gx=gx.contiguous()
                gx=gx.permute(1,0,2,3)
                feature_list=[self.tgcn_net[i](gx[i]) for i in range(seq_len)]
                gcn_feat=torch.cat(feature_list,dim=1)
                return gcn_feat
                
            gcn_feat=[]
            for i in range(batch):
                X=gx[i]
                feature_list=[self.tgcn_net[j](X[j]) for j in range(seq_len)]
                feature=torch.cat(feature_list,dim=1)
                gcn_feat.append(feature)

            gcn_feat=torch.cat(gcn_feat,0)
            return gcn_feat

